import torch  ## 需要导入CANN的环境变量set_env.sh才能生效torch_npu
import torch.nn.functional as F
import torch_npu
import numpy as np
import random


# 设置随机种子
seed = 12
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
## 常量定义
QUANTMAX = 127
QUANTMIN = -128


def process_deq_scale(
        deq_scale: torch.Tensor
        ) -> np.ndarray:
    '''
    处理输入atb MLAPreprocessOperation的 de_scale1 和 de_scale2

    1. 获取deq_scale torch张量数据，转换为numpy数组
    2. 将 float32 的 de_scale1 转换为字节码，用int32重新解释
    3. 转换为int64 torch 张量
    return reinterpret_cast<int32>(
                                deq_scale ) ## reinterpret_cast<int32>
                                            .to(int64)

    原代码中使用本次 golden 数据计算完后，如果本次 type_data = float16，
    则将用完的 self.de_scale1 和 self.de_scale2 做该处理，后续输入 atb 调用。

    对于 atb MLAPreprocessOperation，
    当 data_type = float16，de_scale 为 int64 类型；
    当 data_type = bfloat16，de_scale 为 float32 类型
    '''
    ret = torch.frombuffer(deq_scale.numpy().tobytes(), dtype=torch.int32).to(torch.int64)
    return ret


def round_up(
        val: int, 
        align: int
        ) -> int:
    '''
    辅助函数：向上取整
    '''
    if align == 0:
        return 0
    return -(val // -align) * align

def transdata(
        nd_mat, 
        block_size: tuple = (16, 16)
        ):
    ''' 
    二维数据：ND->NZ
    处理输入 atb MLAPreprocessOperation 的 wdqkv和wuq，这两个矩阵要求是NZ格式，所以输入atb之前需要做ND->NZ
    对于 wuk，可以是ND也可以是NZ。
    若ND，输入 shape = [head_num_q, hidden_size_nope_q_head = 128, hidden_size_wuk_head = 512], 
    若NZ，输入 shape = [head_num_q, 32, 128, 16]
    
    该函数对wdqkv 和 wuq 调用（如果将来需要对比atb）
    '''
    r = round_up(nd_mat.shape[0], block_size[0])
    c = round_up(nd_mat.shape[1], block_size[1])
    r_pad = r - nd_mat.shape[0]
    c_pad = c - nd_mat.shape[1]
    nd_mat = F.pad(nd_mat, ((0, r_pad, 0, c_pad)))
    nz_mat = torch.permute(
        torch.reshape(nd_mat, (r // block_size[0], block_size[0], c // block_size[1], block_size[1])), [2, 0, 1, 3]
    )
    nz_mat = torch.reshape(nz_mat, (nz_mat.shape[0], nz_mat.shape[1] * nz_mat.shape[2], nz_mat.shape[3]))
    return nz_mat

def transdata_3d(
        nd_mat, 
        block_size: tuple = (16, 16)
        ):
    '''
    三维数据的ND->NZ相当于其中各个二维矩阵做ND->NZ
    该函数对wuk调用（如果将来需要对比atb）
    '''
    if nd_mat.ndim != 3:
        raise ValueError("Expected a 3-dimensional input array.")
    B, K, N = nd_mat.shape
    processed_slices = []

    for batch_index in range(B):
        current_slice = nd_mat[batch_index]
        nz_mat = transdata(current_slice, block_size)
        processed_slices.append(nz_mat)
    
    result = torch.stack(processed_slices, axis=0)
    return result


class MLAPO: 

    def __init__(self):
        pass
        
    

    def rotate_half_x_golden(
            self, 
            q_temp: torch.Tensor, 
            head_num: int
            ) -> torch.Tensor:
        '''
        对多头 Q 矩阵的 RoPE 操作：
        将 [tokens, head_num_q * hidden_size_rope_q_head] Q RoPE 部分拆分成 head_num_q 个 [tokens, hidden_size_rope_q_head] 的部分，
        然后对这里的每个 [tokens, hidden_size_rope_q_head] 做 rotate half，
        最后再将旋转后的 head_num_q 个 [tokens, hidden_size_rope_q_head] 拼接返回 [tokens, head_num_q * hidden_size_rope_q_head] 结果。
        '''
        # 拆分成 head_num 个 [n,head_dim] 的二维向量
        q_splits = torch.chunk(q_temp, head_num, dim=1)
        # 对每个 [n,head_dim] 向量的第二维进行分割，并对第二块乘以 -1再拼回到第一块前面
        processed_q_splits = []
        for q_split in q_splits:
            # 分割第二维
            first_half, second_half = torch.chunk(q_split, 2, dim=1)
            # 拼接回 [n,head_dim] 的二维向量
            processed_q_split = torch.cat((-second_half, first_half), dim=1)
            processed_q_splits.append(processed_q_split)
        # 将所有处理后的 [n,head_dim] 向量拼回 [n,head_num*head_dim] 的二维向量
        return torch.cat(processed_q_splits, dim=1)

    def pad_head_num(
            self, 
            head_dim_sin_cos: torch.Tensor, 
            head_num: int
            ) -> torch.Tensor:
        '''
        将 [tokens, hidden_size_rope_q_head]的 cos 和 sin 拷贝拓展到 [tokens, head_num_q * hidden_size_rope_q_head] 的维度
        '''
        return torch.tile(head_dim_sin_cos, (1, head_num))

    def rope_x_golden(
            self, 
            q_rope, 
            cos, 
            sin, 
            # data_type, 
    ) -> torch.Tensor: 
        '''
        输入的 q_rope 为 [tokens, head_num_q, hidden_size_rope_q_head] 维度，
        cos sin 为 [tokens, hidden_size_rope_q_head] 维度，dim = 1 仅对应单头。
        在对 q_rope 做 RoPE 之前，需要将 q_rope 进行 reshape，得到新的 dim = 1 为 head_num_q * hidden_size_rope_q_head，
        再将 cos sin 赋值拓展到多头的维度，
        然后再进行 RoPE 计算。
        注意 RoPE 计算的精度转换。
        '''
        tokens = q_rope.shape[0]
        head_num_q = q_rope.shape[1]
        hidden_size_rope_q_head = q_rope.shape[2]
        pad_cos = self.pad_head_num(cos, head_num_q)
        pad_sin = self.pad_head_num(sin, head_num_q)
        
        data_type = cos.dtype

        q_rope = q_rope.reshape(tokens, head_num_q * hidden_size_rope_q_head)
        q_rope_float32 = q_rope.to(torch.float32)
        pad_cos_float32 = pad_cos.to(torch.float32)
        pad_sin_float32 = pad_sin.to(torch.float32)
        
        res = q_rope_float32 * pad_cos_float32 + self.rotate_half_x_golden(q_rope_float32, head_num_q) * pad_sin_float32
        
        res = res.reshape(tokens, head_num_q, hidden_size_rope_q_head)

        return res.to(data_type)
        # return q_rope * pad_cos + self.rotate_half_x_golden(q_rope, head_num_q) * pad_sin

    def matmul_ein_sum(
            self, 
            mat_a : torch.Tensor, ## [tokens, head_num_q, hidden_size_nope_q_head]
            mat_weight : torch.Tensor, ## [head_num_q, hidden_size_nope_q_head, hidden_size_wuk_head]
            transpose_weight : torch.Tensor, 
            # data_type : torch.dtype, 
    ) -> torch.Tensor:
        '''
        输入 mat_a [tokens, head_num_q, hidden_size_nope_q_head]
        输入 weight [head_num_q, hidden_size_nope_q_head, hidden_size_wuk_head]
        两者矩阵乘法应以 head_num_q 为单位分矩阵，每个矩阵为 [tokens, hidden_size_nope_q_head]。
        因此将 mat_a 进行 (1, 0, 2) 的 permute，
        按照 head_num_q 顺序转为 [head_num_q, token, hidden_size_nope_q_head]，
        再与 weight 进行 dim=1、2 的 matmul，
        即进行 batch_count = head_num_q 的 BatchMatmul。
        最终 BatchMatmul 的结果 [head_num_q, tokens, hidden_size_wuk_head] 做 (1, 0, 2) 的permute，
        结果维度为 [tokens, head_num_q,hidden_size_wuk_head]。
        注意Matmul需要用 float32 计算，结果再转回源类型
        '''
        data_type = mat_weight.dtype
        result = torch.permute(
                    torch.matmul(
                        torch.permute(mat_a, (1, 0, 2)).to(torch.float32), 
                        # torch.permute(mat_a, (1, 0, 2)), 
                        (mat_weight.transpose(1, 2) if transpose_weight else mat_weight).to(torch.float32)
                        # mat_weight.transpose(1, 2) if transpose_weight else mat_weight
                        ),
                    (1, 0, 2),
                )  ## [tokens, head_num_q, hidden_size_wuk_head]
        return result.to(data_type)
    
    
    
    


    def rac_golden(
            self, 
            kv_rac: torch.Tensor, 
            slot_mapping: torch.Tensor, 
            kv_cache: torch.Tensor, 
            # data_type : torch.dtype
            # block_size: int
            ) -> torch.Tensor:
        '''
        将K部分RoPE结果与KV部分RMSNorm结果进行处理，加入cache后返回
        将 tokens 个 kv行 通过查询 slop_mapping 映射到 (block_num, block_size) 个位置去。
        注意算法原地输出到 kv_cache 中
        '''
        block_size = kv_cache.shape[1]
        data_type = kv_cache.dtype
        ## 遍历长度为 tokens 的 slot_mapping 
        ## 第 i 个 kv 对应的槽位号为 slot
        for i, slot in enumerate(slot_mapping):
            if slot < 0:
                continue
            ## block_size 是128、256随机选择
            block_index = slot // block_size
            block_offset = slot % block_size
            token_kv = kv_rac[i].to(data_type)
            kv_cache[block_index][block_offset] = token_kv
        return kv_cache

    # RoPE 旋转半部分（K 专用）
    def rotate_half_golden(
            self, 
            k_temp: torch.Tensor
            ) -> torch.Tensor:
        '''
        对单矩阵 [tokens, hidden_size_rope_k] 的旋转操作
        '''
        first_half, second_half = torch.chunk(k_temp, 2, dim=1)
        processed_k_split = torch.cat((-second_half, first_half), dim=1)
        return processed_k_split


    def rope_golden(
            self, 
            k_rope: torch.Tensor, 
            cos: torch.Tensor,
            sin: torch.Tensor
            ) -> torch.Tensor:
        '''
        输入 tokens 降维切分后，对 K 部分计算 RoPE
        注意 RoPE 计算的精度转换。
        '''
        # src_type = k_rope.dtype
        k_rope_float32 = k_rope.to(torch.float32)
        sin_float32 = sin.to(torch.float32)
        cos_float32 = cos.to(torch.float32)
        res = k_rope_float32 * cos_float32 + self.rotate_half_golden(k_rope_float32) * sin_float32
        return res  # .to(src_type)
        # return k_rope * cos + self.rotate_half_golden(k_rope) * sin

    def rms_norm_golden(
            self, 
            x: torch.Tensor, 
            gamma: torch.Tensor, 
            # hidden_size: int = None, 
            epsilon: float = 1e-5
            ) -> torch.Tensor:
        '''
        RMSNorm 无量化，用于对 输入tokens 经 RMSNormQuant 降维后切分为KV后其中的V部分计算

        gamma * x_float / x的均方根（1方和2均2.5加epsilon3根）
        result = gamma_float * x_float 
                                        /  
                                        sqrt(
                                                sum(
                                                    square(x_float), axis = -1
                                                                                ) ## sum
                                                                                / hidden_size + epsilon 
                                                                                                                    ) ## sqrt
        x_float = x.to(float)
        gamma_float = gamma.to(float)
        '''
        hidden_size = x.shape[-1]
        x_float32 = x.to(torch.float32)
        square_sum = torch.sum(torch.square(x_float32), dim=-1, keepdims=True)
        rms = 1.0 / torch.sqrt(square_sum / hidden_size + epsilon)
        gamma_float32 = gamma.to(torch.float32)
        rms_norm = rms * x_float32 * gamma_float32
        return rms_norm  ##.to(x.dtype)



    def split_rms_norm_rope_reshape_cache_golden(
            self, 
            kv : torch.Tensor, 
            gamma : torch.Tensor, 
            cos : torch.Tensor, 
            sin : torch.Tensor, 
            slot_mapping : torch.Tensor, 
            kv_cache : torch.Tensor, 
            hidden_size_rope_k : int = 64, 
            epsilon : float = 1e-5, 
    ) -> torch.Tensor:
        '''
        对输入 tokens 降维后切分出的 kv 部分再切分：
        对K部分做RoPE，对V部分做RMSNorm，
        reshape后调用rac_golden计算新 kv_cache。
        这部分在 kernel 中可以不用 split 
        '''
        tokens = kv.shape[0]
        hidden_size_kv = kv.shape[-1]
        hidden_size_rms_kv = hidden_size_kv - hidden_size_rope_k

        kv_rms, k_rope = torch.split(
            kv, 
            [hidden_size_rms_kv, hidden_size_rope_k], 
            dim = 1
        )

        rms_norm_output = self.rms_norm_golden(kv_rms, gamma, epsilon)
        rms_norm_reshape = rms_norm_output.reshape(tokens, 1, hidden_size_rms_kv)

        rope_output = self.rope_golden(k_rope, cos, sin)
        rope_reshape = rope_output.reshape(tokens, 1, hidden_size_rope_k)

        kv_rac = torch.cat((rms_norm_reshape, rope_reshape), dim=-1)
        return self.rac_golden(kv_rac, slot_mapping, kv_cache)



    def matmul_w8a8_dequant(
            self, 
            mat_a : torch.Tensor, 
            mat_weight : torch.Tensor, 
            bias: torch.Tensor, 
            de_scale : torch.Tensor, 
            transpose_weight : bool, 
            data_type_out : torch.dtype
    ) -> torch.Tensor: 
        '''
        w8a8计算，结果加偏移再做反量化。
        注意，输入的 mat_a 为 RMSNorm Quant 的结果，元素类型为 int8，
        但 w8a8 dequant 结果应该回到原本类型，
        所以这里需要额外传入的源数据类型。
        TODO 这里应该是int8计算 int32累计，中间不需要转float32
        '''
        matmul_out = torch.matmul(
            mat_a.to(torch.float32), 
            ( mat_weight.transpose(0, 1) if transpose_weight else mat_weight ).to(torch.float32)
        )
        matmul_out_bias = matmul_out + bias
        matmul_out_bias_deq = (matmul_out_bias.to(torch.float32) * de_scale).to(data_type_out)

        return matmul_out_bias_deq

    def rms_norm_quant_golden(
            self, 
            x: torch.Tensor, 
            gamma: torch.Tensor, 
            beta: torch.Tensor,
            quant_scale: torch.Tensor, 
            quant_offset: torch.Tensor, 
            epsilon: float = 1e-5,
            quant_max: int = QUANTMAX, 
            quant_min: int = QUANTMIN
    ) -> torch.Tensor:
        '''
        RMSNorm + Quant 用于，用于 MatMul 输入，
        包括：
        1 输入 tokens 经过 RMSNormQuant 后与 wdqkv 做 matmul降维
        2 输入 tokens 经过 RMSNormQuant 后降维，切分后的 Q 与 wuq 做 matmul 升维

        return  round( 
                    ( RMSNorm(x) + beta ) / quant_scale.to(float) 
                                                                        + quant_offset.to(float) 
                                                                                                ) ## round
                                                                                                    ) 
                                                                                                    .to(fp16)
                                                                                                                .to(int8)
        RMSNorm(x) = gamma_float * input_float / 
                                                    sqrt( 
                                                        sum(
                                                                square(input_float), axis = -1
                                                                                                ) ## sum
                                                                                                    / input.shape[-1] + epsilon ) ## sqrt
        input_float = x.to(float)
        gamma_float = gamma.to(float)
        注意输出结果为 int8 类型
        '''
        # if hidden_size is None:
        hidden_size = x.shape[-1]
        x_float32 = x.to(torch.float32)
        gamma_float32 = gamma.to(torch.float32)
        # 这里dim=-1 一行一行的处理内容
        square_sum = torch.sum(torch.square(x_float32), dim=-1, keepdims=True)
        factor = 1.0 / torch.sqrt(square_sum / hidden_size + epsilon)
        output = x_float32 * factor * gamma_float32

        scale = 1.0 / quant_scale.item()
        offset = quant_offset.item()
        input_scale = torch.tensor(scale, dtype=torch.float32)
        input_offset = torch.tensor(offset, dtype=torch.float32)
        output = (output + beta.to(torch.float32)) * input_scale + input_offset
        # print(output[636][985].item())
        # round->clamp->quant
        # 取整函数（不用torch.round）
        output = torch.where(output >= 0.0, torch.floor(output + 0.5), torch.ceil(output - 0.5))
        # clamp量化到int8的空间中
        output = torch.clamp(output, quant_min, quant_max).to(torch.int8)
        return output


    def mla_preprocess_calc(
            self, 
            input : torch.Tensor, 
            gamma1 : torch.Tensor, 
            beta1 : torch.Tensor, 
            quant_scale1 : torch.Tensor,
            quant_offset1 : torch.Tensor, 
            wdqkv : torch.Tensor, 
            de_scale1 : torch.Tensor, 
            bias1 : torch.Tensor,
            # 
            gamma2 : torch.Tensor, 
            beta2 : torch.Tensor, 
            quant_scale2 : torch.Tensor, 
            quant_offset2 : torch.Tensor,
            wuq : torch.Tensor, 
            de_scale2 : torch.Tensor, 
            bias2 : torch.Tensor, 
            # 
            gamma3 : torch.Tensor,
            cos1 : torch.Tensor, 
            sin1 : torch.Tensor, 
            cos2 : torch.Tensor, 
            sin2 : torch.Tensor, 
            wuk : torch.Tensor, 
            kv_cache : torch.Tensor,
            slot_mapping : torch.Tensor, 
            # 
            hidden_size_wdq : int = 1536, 
            hidden_size_rope_q_head : int = 64, 
            hidden_size_rope_k : int = 64, 
            epsilon : float = 1e-5,  
            transpose_wdqkv : bool = True, 
            transpose_wuq : bool = True, 
            transpose_wuk : bool = False,  
            cache_mode : int = 0, 
            ) -> list:
        """
        MLA Preprocess 全计算流程。

        以main函数传参为例：

        ## 元数据：

        tokens = 1024
        head_num_q = random.choice([32, 128, 16, 64])
        block_size = random.choice([128, 256])
        block_num = 192
        data_type = torch.float16  # 或 torch.bfloat16，只有两个取值
        hidden_size = 7168

        hidden_size_wdqkv = 2112                                                ## 输入 tokens 经过降维后hidden_size
        hidden_size_wdq = 1536                                                  ## 降维后拆分的Q的hidden_size，剩余为KV
        hidden_size_wdkv = hidden_size_wdqkv - hidden_size_wdq                  ## d是降维，u是升维
        hidden_size_rope_k = 64                                                 ## KV传入RoPE的hidden_size（只是K的），剩余为K剩余部分+V
        hidden_size_rms_kv = hidden_size_wdkv - hidden_size_rope_k 
        hidden_size_wuq_head = 192                                              ## 输入tokens经降维+切出Q后升维，其中单头hidden_size
        hidden_size_rope_q_head = hidden_size_rope_k = 64                       ## Q单头传入RoPE的hidden_size
        hidden_size_nope_q_head = hidden_size_wuq_head - hidden_size_rope_q_head
        hidden_size_wuk_head = 512                                              ## 输入tokens经降维+切出Q后升维，其中单头切分为RoPE和NoPE，这里是NoPE升维后hidden_size

        hidden_size_output_kv = hidden_size_rope_k + hidden_size_rms_kv  ## 576
        hidden_size_output_q = hidden_size_rope_q_head + hidden_size_wuk_head  ## 576

        epsilon = 1e-5
        cache_mode = 0
        
        ## 待添加参数 + 功能
        q_rotary_coeff Q旋转系数
        k_rotary_coeff K旋转系数（RoPE的两种实现方式。atb和当前测试都未实现）
        quantMode 指定RMSNorm量化的类型。
                0-PER_TENSOR_QUANT_ASYMM，per_tensor静态非对称量化，默认量化类型。
                1-PER_TOKEN_QUANT_SYMM：per_token 动态对称量化，未实现。
                2-PER_TOKEN_QUANT_ASYMM：per_token 动态非对称量化，未实现。
                3-UNQUANT：不量化，浮点输出，未实现。
                1-3在abt和当前测试都未实现
        cache_mode 指定最终QKV结果输出格式
                - cacheMode 为 0 时，kcache 和 q 均经过拼接后输出。
                - cacheMode 为 1 时，输入输出的 kvCcache 拆分为 krope 和 ctkv，q 拆分为 qrope 和 qnope。
                - cacheMode 为 2 时，krope 和 ctkv 转为 NZ 格式输出，ctkv 和 qnope 经过 per_head 静态对称量化为 int8 类型。
                - cacheMode 为 3 时，krope 和 ctkv 转为 NZ 格式输出。
        
        ## 输出结果：
        qOut 输出 tensor 数据类型data_type｡ cacheMode = 2 时数据类型为 int8
            cacheMode 为 0 : [tokenNum, headNum,576]；
            cacheMode 为 1 : [tokenNum, headNum,64]；
            cacheMode 为 2 : [tokenNum, headNum,512]；
            cacheMode 为 3 : [tokenNum, headNum,512]

        kvCacheOut 输出 tensor  数据类型data_type cacheMode = 2 时数据类型为 int8 ，格式为 NZ 。cacheMode = 3 时，格式为 NZ ｡
            cacheMode 为 0 : [blockNum, blockSize,1, 576]；
            cacheMode 为 1 : [blockNum, blockSize,1, 512]；
            cacheMode 为 2 : [blockNum, headNum_512/32,block_size, 32]；
            cacheMode 为 3 : [blockNum, headNum_512/16,block_size, 16]

        kvCacheRopeOut 为 0 时输出此 tensor ｡数据类型 = data_type｡ cacheMode = 2 时数据格式为 NZ ｡ cacheMode = 3 时，格式为 NZ ｡
            cacheMode 不为 0 : [tokenNum,headNum,64]；
            cacheMode 为 1 : [blockNum, blockSize,1, 64]；
            cacheMode 为 2 : [blockNum, headNum_64 / 16 ,block_size, 16]；
            cacheMode 为 3 : [blockNum, headNum_64 / 16 ,block_size, 16]
                
        ## 入参：

        :param input : 输入 tokens（包含QKV）                      shape = (tokens, hidden_size), type = data_type
        :param gamma1 : 对输入 tokens 做 RMSNormQuant 参数          shape = (hidden_size), type = data_type
        :param beta1 : 对输入 tokens 做 RMSNormQuant 参数           shape = (hidden_size), type = data_type
        :param quant_scale1 : 对输入 tokens 做 RMSNormQuant 参数    shape = 1, type = data_type
        :param quant_offset1 : 对输入 tokens 做 RMSNormQuant 参数   shape = 1, type = int8
        :param wdqkv : 输入 tokens 的降维矩阵用以matmul             shape = (hidden_size_wdqkv, hidden_size), type = int8
        :param de_scale1 : 输入 tokens 降维后缩放参数               shape = (hidden_size_wdqkv), type = float32
        :param bias1 : 输入 tokens 降维后偏移                       shape = (1, hidden_size_wdqkv), type = int32
        
        :param gamma2 : Q部分 RMSNormQuant 参数                     shape = (hidden_size_wdq), type = data_type
        :param beta2 : Q部分 RMSNormQuant 参数                      shape = (hidden_size_wdq), type = data_type
        :param quant_scale2 : Q部分 RMSNormQuant 参数               shape = 1, type = data_type
        :param quant_offset2 : Q部分 RMSNormQuant 参数              shape = 1, type = int8
        :param wuq : Q部分的升维矩阵用以matmul                      shape = (head_num_q * hidden_size_wuq_head, hidden_size_wdq), type = int8
        :param de_scale2 : Q部分升维后缩放参数                      shape = (head_num_q * hidden_size_wuq_head), type = float32
        :param bias2 : Q部分升维后偏移                              shape = (1, head_num_q * hidden_size_wuq_head)， type = int32
        
        :param gamma3 : KV 部分 RMSNorm 参数                         shape = (hidden_size_rms_kv), type = data_type
        :param cos1 : K 部分 RoPE 参数                              shape = (tokens, hidden_size_rope_k), type = data_type
        :param sin1 : K 部分 RoPE 参数                              shape = (tokens, hidden_size_rope_k), type = data_type
        :param cos2 : Q 部分单头64维度 RoPE 参数                    数据同sin2
        :param sin2 : Q 部分单头64维度 RoPE 参数                    数据同sin1
        :param wuk : Q 部分单头128维度 升维矩阵用以matmul           shape = (head_num_q, hidden_size_nope_q_head, hidden_size_wuk_head), type = data_type
        :param kv_cache : K（和V）缓存？                           shape = (block_num, block_size, 1, hidden_size_output_kv), type = data_type 
        :param slot_mapping: 槽位映射 (tokens)                      shape = (tokens), type = int32

        :param hidden_size_wdq : 降维后的tokens切分k维度 1536
        :param hidden_size_rope_q_head : q 单头 RoPE 隐藏维度 64
        :param hidden_size_rope_k : k RoPE 隐藏维度 64
        :param epsilon : RMS Norm epsilon
        :transpose_wdqkv : wdqkv 是否转置，默认 True
        :transpose_wuq : wuq 是否转置，默认 True
        :transpose_wuk : wuk 是否转置，默认 False
        :param cache_mode : 缓存模式 (0: 无 split; 1: split Q/KV)
        
        ## 结果
        :return : [q_out, kv_out] 或 split 版本

        TODO 
        与官方torch文件不同：
        1. 官方的rope部分（包括Q和KV）都是将输入数据与cos和sin直接用data_type计算，没有将三者转为float
        2. 官方的w8a8+dequant结果最终转回data_type，没有维持float32进行下一步计算。
        """

        tokens = input.shape[0]
        head_num_q = wuk.shape[0]
        data_type = input.dtype
        
        # RMSNorm + Quant 处理输入 tokens
        wdqkv_type = wdqkv.dtype
        rms_quant_out1 = None
        if wdqkv_type == torch.bfloat16: 
            rms_quant_out1 = input
        else: 
            rms_quant_out1 = self.rms_norm_quant_golden(
                input, 
                gamma1, 
                beta1, 
                quant_scale1, 
                quant_offset1, 
                epsilon, 
                )
        # print("rms_quant_out1 type: ", rms_quant_out1.dtype)
        # Matmul bias DeQuant 对 RMSNormQuant处理的tokens降维
        # return rms_quant_out1.float()

        mm1_out = self.matmul_w8a8_dequant(
            rms_quant_out1, 
            wdqkv, 
            bias1, 
            de_scale1, 
            transpose_wdqkv, 
            data_type
        )
        # print("mm1_out type: ", mm1_out.dtype)
        # return [mm1_out]

        # Split 经过RMSNormQuant处理后降维的tokens [2112(qkv)] -> [576(kv), 1536 (q)]
        hidden_size_wdkv = mm1_out.shape[-1] - hidden_size_wdq
        mm1_split1, mm1_split2 = torch.split(
            mm1_out, 
            [hidden_size_wdkv, hidden_size_wdq], 
            dim=1
            )

        # RMSNorm + Quant 处理经过 RMSNormQuant 处理后再降维后 tokens 中的 q 部分
        wuq_type = wuq.dtype
        rms_quant_out2 = None
        if wuq_type == torch.bfloat16: 
            rms_quant_out2 = mm1_split2
        else: 
            rms_quant_out2 = self.rms_norm_quant_golden(
                mm1_split2, 
                gamma2, 
                beta2, 
                quant_scale2, 
                quant_offset2, 
                epsilon,  
                # hidden_size_wdq
                )
        # return rms_quant_out2.float()
        # print("rms_quant_out2 type: ", rms_quant_out2.dtype)

        # MatMul bias DeQuant  对 经过 RMSNormQuant 处理后降维后 tokens 中的 q 部分做 RMSNormQuant 后 升维，总维度为 head_num_q * wuq.shape[-1]

        mm2_out = self.matmul_w8a8_dequant(
            rms_quant_out2, 
            wuq, 
            bias2, 
            de_scale2, 
            transpose_wuq, 
            data_type
        )
        # return [mm2_out]
        # print("mm2_out type: ", mm2_out.dtype)

        # split 经过 RMSNormQuant 处理后降维后 tokens 中的 kv 部分，分别做 RMSNorm 和 RoPE，最后 Reshape Cache 处理
        kv_out = self.split_rms_norm_rope_reshape_cache_golden(
            mm1_split1, 
            gamma3, 
            cos1, 
            sin1, 
            slot_mapping, 
            kv_cache, 
            hidden_size_rope_k, 
            epsilon, 
            # data_type, 
        )
        # return [mm2_out, kv_out]
        # print("kv_out type: ", kv_out.dtype)

        # Split 经过 RMSNormQuant 处理后降维后 tokens 中的 q部分做 RMSNormQuant 后升维的单头，[192 q] -> [128 (q_nope), 64 (q_rope)]
        hidden_size_wuq_head = mm2_out.shape[-1] // head_num_q
        mm2_out = mm2_out.reshape(tokens, head_num_q, hidden_size_wuq_head)
        hidden_size_nope_q_head = hidden_size_wuq_head - hidden_size_rope_q_head
        mm2_split1, mm2_split2 = torch.split(
            mm2_out, 
            [hidden_size_nope_q_head, hidden_size_rope_q_head], 
            dim=2
            )

        # MatmulEinSum，对 Q 切分后的单头中的 NoPE 部分升维，输出 [tokens, head_num_q, hidden_size_wuk_head]
        bmm_out = self.matmul_ein_sum(
            mm2_split1, 
            wuk, 
            transpose_wuk, 
            # data_type, 
        )
        # print("bmm_out type: ", bmm_out.dtype)
        # return [bmm_out]

        ## q rope 部分，输出 [tokens, head_num_q, hidden_size_rope_q_head]
        rope_q_out = self.rope_x_golden(
            mm2_split2, 
            cos2, 
            sin2, 
            # data_type, 
        )
        # print("rope_q_out: ", rope_q_out.dtype)
        
        # return [kv_out, rope_q_out]

        q_out = torch.cat((bmm_out, rope_q_out), dim = 2)

        # 根据 cache_mode 返回输出
        if cache_mode == 1:
            return [
                q_out[..., 0:512], 
                kv_out[..., 0:512],
                q_out[..., 512:576], 
                kv_out[..., 512:576]
            ]
        else:
            return [q_out, kv_out]



